import h5py
import numpy as np
from skimage.morphology import ball, disk, dilation, binary_erosion, remove_small_objects, erosion, closing, \
    reconstruction, binary_closing, remove_small_holes
from skimage.measure import label, regionprops
from skimage.morphology import binary_dilation, binary_opening
from skimage.segmentation import clear_border
from skimage import data
from scipy import ndimage as ndi


def fix_contour_slice(bi_mask, spacing):
    spacing_y = spacing[0]
    spacing_x = spacing[1]
    fixed_bi_mask = np.zeros(bi_mask.shape, dtype=np.bool)
    labeled_bi_mask = label(bi_mask)
    regions = regionprops(labeled_bi_mask)
    disk_12 = disk(15, dtype=np.bool)
    for region in regions:
        mask_now = bi_mask == region.label
        mask_now = binary_dilation(mask_now, disk_12)
        mask_now = remove_small_holes(mask_now, min_size=20/spacing_x*20/spacing_y)
        mask_now = binary_erosion(mask_now, disk_12)
        fixed_bi_mask = fixed_bi_mask + mask_now
    return fixed_bi_mask


def lung_segmentation_scan(vol, spacing=(0.8, 0.8, 1)):
    """
    segment Lung volume
    :param 
    vol: Lung volume, original HU values are required
    spacing: volume spacing
    :return: binary LungMask
    """
    if type(vol) is not np.ndarray:
        vol = np.array(vol)
    bi_vol = vol < -400
    for i in range(vol.shape[2]):
        bi_vol[:, :, i] = clear_border(bi_vol[:, :, i])
    label_vol = label(bi_vol, neighbors=8)
    regions = regionprops(label_vol)
    max_region = dict(index=0, area=0)
    for i, region in enumerate(regions):
        if region.area > max_region['area']:
            max_region['index'] = i
            max_region['area'] = region.area
    bi_vol = label_vol == regions[max_region['index']].label

    for i in range(vol.shape[2]):
        bi_vol[:, :, i] = fix_contour_slice(bi_vol[:, :, i], spacing)

    bi_vol = binary_closing(bi_vol, ball(3))
    bi_vol = binary_dilation(bi_vol, ball(1))
    return bi_vol


if __name__ == '__main__':
    LIDC_vol_file = h5py.File('/data_4t/Kaggle/backup/lidc/vol.hdf5', 'r')
    vol = LIDC_vol_file['LIDC-IDRI-0001_01'][:]
    vol = lung_segmentation_scan(vol)
